package cn.edu.cqu.fredyvia;

import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.ml.evaluation.RegressionEvaluator;
import org.apache.spark.ml.recommendation.ALS;
import org.apache.spark.ml.recommendation.ALSModel;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.SparkSession;
import org.apache.spark.sql.functions;
import org.apache.spark.sql.Row;

public class Main {
  public static void main(String[] args) {
    SparkSession spark = SparkSession.builder().appName("Simple Application").getOrCreate();
    Dataset<Row> ratings = spark.read().option("header", true).option("inferSchema", true)
        .csv("hdfs://namenode.fredyvia.asia:9000/ratings.csv");
    Dataset<Row>[] splits = ratings.randomSplit(new double[] { 0.8, 0.2 });
    Dataset<Row> training = splits[0];
    Dataset<Row> test = splits[1];

    // Build the recommendation model using ALS on the training data
    ALS als = new ALS()
        .setMaxIter(5)
        .setRegParam(0.01)
        .setUserCol("userId")
        .setItemCol("movieId")
        .setRatingCol("rating");
    ALSModel model = als.fit(training);

    // Evaluate the model by computing the RMSE on the test data
    // Note we set cold start strategy to 'drop' to ensure we don't get NaN
    // evaluation metrics
    model.setColdStartStrategy("drop");
    Dataset<Row> predictions = model.transform(test);

    RegressionEvaluator evaluator = new RegressionEvaluator()
        .setMetricName("rmse")
        .setLabelCol("rating")
        .setPredictionCol("prediction");
    double rmse = evaluator.evaluate(predictions);
    System.out.println("Root-mean-square error = " + rmse);

    // Generate top 10 movie recommendations for each user
    Dataset<Row> userRecs = model.recommendForAllUsers(10);
    userRecs.withColumn(
        "recommendations", functions.to_json(userRecs.col("recommendations")))
        .orderBy("movieId");
    // Generate top 10 user recommendations for each movie
    Dataset<Row> movieRecs = model.recommendForAllItems(10);
    movieRecs.withColumn(
        "recommendations", functions.to_json(userRecs.col("recommendations")))
        .orderBy("movieId");
    userRecs.write().option("header", true).option("inferSchema", true)
        .csv("hdfs://namenode.fredyvia.asia:9000/users_java");
    movieRecs.write().option("header", true).option("inferSchema", true)
        .csv("hdfs://namenode.fredyvia.asia:9000/movies_java");
  }
}